"""
股票实时数据抓取模块
"""
import datetime
import time
import pandas as pd
from loguru import logger
import akshare as ak

from database.db_engine import SessionLocal
from models.stock_models import StockBasicInfo, StockRealtimeData

def fetch_realtime_data():
    """
    获取所有A股股票的实时交易数据（包括上证、深证、创业板）
    
    Returns:
        pandas.DataFrame: 实时数据
    """
    try:
        logger.info("开始获取实时股票交易数据...")
        
        # 使用akshare获取实时行情数据
        realtime_df = ak.stock_zh_a_spot_em()
        
        # 确保代码列为字符串类型
        realtime_df['代码'] = realtime_df['代码'].astype(str)
        
        # 根据股票代码添加交易所标识
        def get_exchange_code(stock_code):
            if stock_code.startswith('6'):
                return '.SH'  # 上证
            elif stock_code.startswith('3'):
                return '.CY'  # 创业板
            else:
                return '.SZ'  # 深证
        
        # 转换列名
        renamed_df = realtime_df.rename(
            columns={
                '代码': 'stock_code',
                '名称': 'name',
                '最新价': 'price',
                '涨跌额': 'change',
                '涨跌幅': 'pct_change',
                '成交量': 'volume',
                '成交额': 'amount'
            }
        )
        
        # 添加交易所代码和时间
        renamed_df['ts_code'] = renamed_df['stock_code'].apply(lambda x: x + get_exchange_code(x))
        renamed_df['trade_time'] = datetime.datetime.now()
        
        # 转换数据类型
        renamed_df['price'] = pd.to_numeric(renamed_df['price'], errors='coerce')
        renamed_df['change'] = pd.to_numeric(renamed_df['change'], errors='coerce')
        renamed_df['pct_change'] = renamed_df['pct_change'].astype(str).str.replace('%', '').astype(float)
        renamed_df['volume'] = pd.to_numeric(renamed_df['volume'], errors='coerce')
        renamed_df['amount'] = pd.to_numeric(renamed_df['amount'], errors='coerce')
        
        # 提取需要的字段，确保字段顺序与数据库模型一致
        result_df = renamed_df[[
            'ts_code',
            'stock_code',
            'trade_time',
            'price',
            'change',
            'pct_change',
            'volume',
            'amount'
        ]]
        
        logger.info(f"成功获取实时股票交易数据，共 {len(result_df)} 条记录")
        logger.info(f"数据统计：上证股票 {len(result_df[result_df['ts_code'].str.endswith('.SH')])} 条")
        logger.info(f"数据统计：深证股票 {len(result_df[result_df['ts_code'].str.endswith('.SZ')])} 条")
        logger.info(f"数据统计：创业板股票 {len(result_df[result_df['ts_code'].str.endswith('.CY')])} 条")
        return result_df
    
    except Exception as e:
        logger.error(f"获取实时股票交易数据失败: {e}")
        logger.exception("详细错误信息:")
        return pd.DataFrame()

def update_realtime_data():
    """
    更新实时股票交易数据到数据库
    
    Returns:
        bool: 操作是否成功
    """
    try:
        # 获取实时数据
        realtime_df = fetch_realtime_data()
        
        if realtime_df.empty:
            logger.warning("未获取到实时股票交易数据，跳过更新")
            return False
            
        # 获取数据库会话
        db = SessionLocal()
        
        try:
            # 批量插入数据
            for _, row in realtime_df.iterrows():
                # 将数据转换为字典，并处理nan值
                realtime_dict = row.where(pd.notnull(row), None).to_dict()
                
                # 由于是实时数据，直接创建新记录
                new_data = StockRealtimeData(**realtime_dict)
                db.add(new_data)
            
            # 提交事务
            db.commit()
            logger.info(f"成功更新实时股票交易数据，共 {len(realtime_df)} 条记录")
            return True
            
        except Exception as e:
            db.rollback()
            logger.error(f"更新实时股票交易数据到数据库失败: {e}")
            return False
        finally:
            db.close()
            
    except Exception as e:
        logger.error(f"更新实时股票交易数据过程异常: {e}")
        return False

def fetch_and_save_realtime_data():
    """
    获取并保存实时股票交易数据的主函数
    
    Returns:
        bool: 操作是否成功
    """
    logger.info("开始执行实时数据抓取任务...")
    
    # 交易时间判断
    now = datetime.datetime.now()
    current_time = now.time()
    
    # 只在交易时间内执行（9:30-11:30, 13:00-15:00，工作日）
    is_trading_time = (
        (current_time >= datetime.time(9, 30) and current_time <= datetime.time(11, 30)) or
        (current_time >= datetime.time(13, 0) and current_time <= datetime.time(15, 0))
    )
    is_weekday = now.weekday() < 5  # 0-4 对应周一至周五
    
    if is_trading_time and is_weekday:
        logger.info("当前为交易时段，开始抓取实时数据...")
        return update_realtime_data()
    else:
        logger.info("当前非交易时段，跳过实时数据抓取")
        return False

if __name__ == "__main__":
    # 测试代码
    fetch_and_save_realtime_data() 